import os
from torch.backends import cudnn
from utils.logger import setup_logger
from datasets import make_dataloader
from model.DouZi import make_model
from solver import make_optimizer, WarmupMultiStepLR
from loss import make_loss_RMGL
# from processor import do_train
from processor.Douzi_processor import do_train
import random
import torch
import numpy as np
import os
import argparse
from config import cfg

import logging
import torch.nn as nn

try:
    from torch.cuda import amp
except ImportError:
    pass

from model.DouZi import feat_channel_branch


def set_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = True


if __name__ == '__main__':

    parser = argparse.ArgumentParser(description="ReID Baseline Training")
    parser.add_argument(
        "--config_file", default="", help="path to config file", type=str
    )
    parser.add_argument("opts", help="Modify config options using the command-line", default=None,
                        nargs=argparse.REMAINDER)

    args = parser.parse_args()

    if args.config_file != "":
        cfg.merge_from_file(args.config_file)
    cfg.merge_from_list(args.opts)
    cfg.freeze()

    set_seed(cfg.SOLVER.SEED)

    output_dir = cfg.OUTPUT_DIR
    if output_dir and not os.path.exists(output_dir):
        os.makedirs(output_dir)

    logger = setup_logger("reid_baseline", output_dir, if_train=True)
    logger.info("Saving model in the path :{}".format(cfg.OUTPUT_DIR))
    logger.info(args)

    if args.config_file != "":
        logger.info("Loaded configuration file {}".format(args.config_file))
        with open(args.config_file, 'r') as cf:
            config_str = "\n" + cf.read()
            logger.info(config_str)
    logger.info("Running with config:\n{}".format(cfg))

    os.environ['CUDA_VISIBLE_DEVICES'] = cfg.MODEL.DEVICE_ID
    train_loader, val_loader_green, val_loader_normal, num_query_green, num_query_normal, num_classes = make_dataloader(cfg)

    # if cfg.MODEL.PRETRAIN_CHOICE == 'finetune':
    #     model = make_model(cfg, num_class=num_classes)  # 在已经训练参数基础上进行训练，断点续训
    #     # model.load_param_finetune(cfg.MODEL.PRETRAIN_PATH)  # 只加载部分参数
    #     # model.load_param(cfg.MODEL.PRETRAIN_PATH)  # 加载全部参数
    # else:
    #     model = make_model(cfg, num_class=num_classes)
    #     start_epoch = 0

    logger = logging.getLogger("reid_baseline.train")
    "1. 模型构建"
    model = make_model(cfg, num_class=num_classes)
    "构建loss计算方法，返回总loss和center loss值"
    loss_func, center_criterion = make_loss_RMGL(cfg, num_classes=num_classes, feat_dim=feat_channel_branch)
    "构建优化器，center 固定使用SGD"
    optimizer, optimizer_center = make_optimizer(cfg, model, center_criterion)
    "预定义amp scaler"
    scaler = None
    if cfg.SOLVER.FP16:
        logger.info("use FP16 for train")
        "构建amp scaler"
        scaler = amp.GradScaler()

    "3. 加载预训练模型各参数"
    start_epoch = 0
    if cfg.MODEL.PRETRAIN_CHOICE == 'finetune':
        print('Loading pretrained model from {} for finetuning......'.format(cfg.MODEL.PRETRAIN_PATH))
        checkpoint = torch.load(cfg.MODEL.PRETRAIN_PATH)  # 参数先存入CPU
        model.load_param_finetune_from_dict(checkpoint['model_state_dict'])
        # model.load_state_dict(checkpoint['model_state_dict'])
        start_epoch = checkpoint['epoch']
        # start_epoch = 0  # 使学习率从头开始变换，注释后则按照加载模型的位置开始
        print('start epoch:{}'.format(start_epoch))

    scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS, cfg.SOLVER.GAMMA,
                                  cfg.SOLVER.WARMUP_FACTOR,
                                  cfg.SOLVER.WARMUP_EPOCHS, cfg.SOLVER.WARMUP_METHOD,
                                  last_epoch=start_epoch - 1)

    "2. to gpu 否则预训练模型无法加载"
    device = "cuda"
    if device:
        model.to(device)  # 参数转移到GPU
    "加载optimizer"
    # if cfg.MODEL.PRETRAIN_CHOICE == 'finetune':
    #     optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    #     optimizer_center.load_state_dict(checkpoint['optimizer_center_state_dict'])
    #     if cfg.SOLVER.FP16:
    #         scaler.load_state_dict(checkpoint['amp'])

    do_train(
        cfg,
        model,
        center_criterion,
        train_loader,
        val_loader_green,  # 包含了query 和 gallery
        val_loader_normal,
        optimizer,
        optimizer_center,
        scheduler,  # modify for using self trained model
        loss_func,
        num_query_green,  # 验证集的query类别数量
        num_query_normal,
        start_epoch,
        scaler=scaler
    )
    # "2. to gpu 否则预训练模型无法加载"
    # device = "cuda"
    # if device:
    #     model.to(device)
    #     if torch.cuda.device_count() > 1:
    #         print('Using {} GPUs for training'.format(torch.cuda.device_count()))
    #         model = nn.DataParallel(model)
    #     else:
    #         if cfg.SOLVER.FP16:
    #             logger.info("use FP16 for train")
    #             "构建amp scaler"
    #             scaler = amp.GradScaler()
    #             # model, optimizer = amp.initialize(model, optimizer, opt_level='O1')
